#include <Dynamics/Pdm/pdm_cuda.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <Utilities/cuda_utilities.h>
#include <Utilities/SpaceHash/space_hash_cuda.h>
#include <cuda_block_size_set.h>
using namespace std;

namespace PhysLeo {

template<typename T>
PdmCuda<T>::PdmCuda():PdmCpu<T>()
{
}

template<typename T>
void PdmCuda<T>::allocateResource(SimData<T>& sim_data)
{
    //buffer data
    this->ptr_rest_pos_ = sim_data.ptr_rest_pos_;
    this->ptr_pos_ = sim_data.ptr_pos_;
    this->ptr_force_ = sim_data.ptr_force_;
    this->ptr_vert_volume_ = sim_data.ptr_vert_volume_;
    this->ptr_vert_material_ = sim_data.ptr_vert_material_;

    //property
    this->horizon_radius_ = sim_data.horizon_radius_;
    this->homogeneous_ = sim_data.homogeneous_;
    this->pdm_material_type_ = sim_data.pdm_material_type_;
    this->pdm_invertible_ = sim_data.pdm_invertible_;
    this->pdm_weight_type_ = sim_data.pdm_weight_type_;

    int vert_num = this->ptr_vert_volume_->size();

    //allocate space for neighbor list
    this->ptr_neighbor_list_ = make_shared<BufferData<NeighborList>>(vert_num);
    this->ptr_neighbor_list_->allocateGpu();

    SpaceHashCuda<T> space_hash;
    auto bbox = space_hash.bbox(this->ptr_rest_pos_);
    space_hash.setSpace(this->horizon_radius_, bbox[0], bbox[1], 2 * SPACE_HASH_MAX_NEIGHBOR);
    space_hash.queryNeighbor(this->ptr_rest_pos_, this->ptr_neighbor_list_, this->horizon_radius_, this->horizon_radius_*0.1f);

#ifndef NDEBUG
    this->ptr_neighbor_list_->gpu2Cpu();
#endif

    //allocate space for weight
    this->ptr_weight_ = make_shared<BufferData<PdmWeight<T>>>(vert_num);
    this->ptr_weight_->allocateGpu();
    initWeights();

    //allocate space for volume dilation
    this->ptr_volume_dilation_ = make_shared<BufferData<T>>(vert_num);
    this->ptr_volume_dilation_->allocateGpu();
}

template<typename T>
__global__ static void K_PdmCuda_deformEnergy(glm::tvec3<T>* rest_pos_ptr, glm::tvec3<T>* pos_ptr, T* energy_ptr, NeighborList* neighborlist_ptr, PdmWeight<T>* weight_ptr, T* volume_dilation_ptr, T* material_ptr, T* volume_ptr, bool homogeneous, PdmMaterialType pdm_material_type, int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    const auto vert_id = pid;

    const auto material_param_num = PdmMaterial<T>::numMaterialParams(pdm_material_type);
    material_ptr = material_ptr + (homogeneous ? 0 : vert_id * material_param_num);

    const int neighbor_size = neighborlist_ptr[vert_id].size_;
    T energy = 0.0f;
    for (auto i = 0; i < neighbor_size; ++i)
    {
        auto dst_index = neighborlist_ptr[vert_id].ids_[i];
        glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
        glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];

        energy += volume_ptr[vert_id] * weight_ptr[vert_id][i] * PdmMaterial<T>::energy(origin, deform, volume_dilation_ptr[vert_id], material_ptr, pdm_material_type);
    }
    energy_ptr[vert_id] = energy;
}

template<typename T>
T PdmCuda<T>::deformEnergy(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    volumeDilation(ptr_pos);

    auto rest_pos_ptr = this->ptr_rest_pos_->ptrGpu();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto neighborlist_ptr = this->ptr_neighbor_list_->ptrGpu();
    auto weight_ptr = this->ptr_weight_->ptrGpu();
    auto volume_dilation_ptr = this->ptr_volume_dilation_->ptrGpu();
    auto material_ptr = this->ptr_vert_material_->ptrGpu();
    auto volume_ptr = this->ptr_vert_volume_->ptrGpu();

    const int vert_num = this->ptr_rest_pos_->size();
    thrust::device_vector<T> energy(vert_num);
    auto energy_ptr = energy.data().get();

    K_PdmCuda_deformEnergy<T> << <numBlocks(vert_num), BLOCK_SIZE >> > (rest_pos_ptr, pos_ptr, energy_ptr , neighborlist_ptr, weight_ptr, volume_dilation_ptr, material_ptr, volume_ptr, this->homogeneous_, this->pdm_material_type_, vert_num);
    synchronCheck;

    auto total_energy = thrust::reduce(energy.begin(), energy.end(), 0.0f, thrust::plus<T>());
    return total_energy;
}

template<typename T>
__global__ static void K_PdmCuda_internalForce(glm::tvec3<T>* rest_pos_ptr, glm::tvec3<T>* pos_ptr, glm::tvec3<T>* force_ptr, NeighborList* neighborlist_ptr, PdmWeight<T>* weight_ptr, T* volume_dilation_ptr, T* material_ptr, T* volume_ptr, bool homogeneous, PdmMaterialType pdm_material_type, int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    const auto vert_id = pid;

#ifndef NDEBUG
    if(pid==0)printf("internal_force before(%f,%f,%f)\n ",force_ptr[vert_id][0],force_ptr[vert_id][1],force_ptr[vert_id][2]);
#endif

    const auto material_param_num = PdmMaterial<T>::numMaterialParams(pdm_material_type);
    material_ptr = material_ptr + (homogeneous ? 0 : vert_id * material_param_num);

    const int neighbor_size = neighborlist_ptr[vert_id].size_;
    glm::tvec3<T> center_force(0.0f);
    for (auto i = 0; i < neighbor_size; ++i)
    {
        auto dst_index = neighborlist_ptr[vert_id].ids_[i];
        glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
        glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];

        glm::tvec3<T> force = volume_ptr[vert_id] * weight_ptr[vert_id][i] * PdmMaterial<T>::force(origin, deform, volume_dilation_ptr[vert_id], material_ptr, pdm_material_type);
        center_force += force;

        atomicAdd(&(force_ptr[dst_index][0]), -force[0]);
        atomicAdd(&(force_ptr[dst_index][1]), -force[1]);
        atomicAdd(&(force_ptr[dst_index][2]), -force[2]);
    }
    atomicAdd(&(force_ptr[vert_id][0]), center_force[0]);
    atomicAdd(&(force_ptr[vert_id][1]), center_force[1]);
    atomicAdd(&(force_ptr[vert_id][2]), center_force[2]);

#ifndef NDEBUG
    if(pid==0)printf("internal force after(%f,%f,%f)\n",force_ptr[vert_id][0],force_ptr[vert_id][1],force_ptr[vert_id][2]);
#endif

    //delete vert_material;
}

template<typename T>
void PdmCuda<T>::internalForce(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos, std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_force)
{
    volumeDilation(ptr_pos);

    auto rest_pos_ptr = this->ptr_rest_pos_->ptrGpu();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto force_ptr = ptr_force->ptrGpu();
    auto neighborlist_ptr = this->ptr_neighbor_list_->ptrGpu();
    auto weight_ptr = this->ptr_weight_->ptrGpu();
    auto volume_dilation_ptr = this->ptr_volume_dilation_->ptrGpu();
    auto material_ptr = this->ptr_vert_material_->ptrGpu();
    auto volume_ptr = this->ptr_vert_volume_->ptrGpu();

    const int vert_num = this->ptr_rest_pos_->size();

    K_PdmCuda_internalForce<T> << <numBlocks(vert_num), BLOCK_SIZE >> > (rest_pos_ptr, pos_ptr, force_ptr, neighborlist_ptr, weight_ptr, volume_dilation_ptr, material_ptr, volume_ptr, this->homogeneous_, this->pdm_material_type_, vert_num);
}

template<typename T>
__global__ static void K_PdmCuda_volumeDilation(glm::tvec3<T>* rest_pos_ptr, glm::tvec3<T>* pos_ptr, NeighborList* neighborlist_ptr, PdmWeight<T>* weight_ptr, T* volume_dilation_ptr, int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    const auto vert_id = pid;

    const int neighbor_size = neighborlist_ptr[vert_id].size_;
    T theta = 0.0f;
    for (auto i = 0; i < neighbor_size; ++i)
    {
        auto dst_index = neighborlist_ptr[vert_id].ids_[i];
        glm::tvec3<T> origin = rest_pos_ptr[dst_index] - rest_pos_ptr[vert_id];
        glm::tvec3<T> deform = pos_ptr[dst_index] - pos_ptr[vert_id];

        auto s = glm::length(deform) / glm::length(origin);
        theta += weight_ptr[vert_id][i] * s;
    }
    volume_dilation_ptr[vert_id] = theta;
}

template<typename T>
void PdmCuda<T>::volumeDilation(std::shared_ptr<BufferData<glm::tvec3<T>>> ptr_pos)
{
    auto rest_pos_ptr = this->ptr_rest_pos_->ptrGpu();
    auto pos_ptr = ptr_pos->ptrGpu();
    auto neighborlist_ptr = this->ptr_neighbor_list_->ptrGpu();
    auto weight_ptr = this->ptr_weight_->ptrGpu();
    auto volume_dilation_ptr = this->ptr_volume_dilation_->ptrGpu();

    const int vert_num = this->ptr_rest_pos_->size();

    K_PdmCuda_volumeDilation<T> << <numBlocks(vert_num), BLOCK_SIZE >> > (rest_pos_ptr, pos_ptr, neighborlist_ptr, weight_ptr, volume_dilation_ptr, vert_num);
    synchronCheck;
}

template<typename T>
__global__ static void K_PdmCuda_initWeights(glm::tvec3<T> *rest_pos_ptr, NeighborList* neighborlist_ptr, T* vert_volume_ptr, PdmWeight<T>* weight_ptr, PdmWeightType pdm_weight_type, T horizon_radius, int size)
{
    int pid = threadIdx.x + (blockIdx.x * blockDim.x);
    if (pid >= size) return;
    const auto i = pid;

    const auto neighbor_size = neighborlist_ptr[i].size_;
    T total_weight = 0.0f;
    for (auto j = 0; j<neighbor_size; ++j)
    {
        auto dst_index = neighborlist_ptr[i].ids_[j];
        auto original_length = glm::length(rest_pos_ptr[i] - rest_pos_ptr[dst_index]);
        switch (pdm_weight_type)
        {
        case PdmWeightType::AVERAGE:
            weight_ptr[i][j] = 1.0f;
            break;
        case PdmWeightType::CENTER:
            weight_ptr[i][j] = horizon_radius / (horizon_radius + original_length);
            break;
        case PdmWeightType::OUTER:
            weight_ptr[i][j] = (horizon_radius + original_length) / horizon_radius;
            break;
        default:
            break;
        }
        weight_ptr[i][j] *= vert_volume_ptr[dst_index];
        total_weight += weight_ptr[i][j];
    }
    for (auto j = 0; j<neighbor_size; ++j)
    {
        //normalization
        weight_ptr[i][j] /= total_weight;
    }
}

template<typename T>
void PdmCuda<T>::initWeights()
{
    auto rest_pos_ptr = this->ptr_rest_pos_->ptrGpu();
    auto neighborlist_ptr = this->ptr_neighbor_list_->ptrGpu();
    auto vert_volume_ptr = this->ptr_vert_volume_->ptrGpu();
    auto weight_ptr = this->ptr_weight_->ptrGpu();

    const int vert_num = this->ptr_rest_pos_->size();
    K_PdmCuda_initWeights<T> << <numBlocks(vert_num), BLOCK_SIZE >> > (rest_pos_ptr, neighborlist_ptr, vert_volume_ptr, weight_ptr, this->pdm_weight_type_, this->horizon_radius_, vert_num);
    synchronCheck;
}

template class PdmCuda<float>;
template class PdmCuda<double>;

}